import zmq
import numpy as np
from funasr import AutoModel
from dataclasses import dataclass
from typing import Dict, List, Optional
from collections import deque
import threading
import time
import logging
from logging.handlers import RotatingFileHandler
import os


# 配置日志
def setup_logging():
    log_dir = "../log"
    # log_dir = "/root/app"
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)

    # 创建日志记录器
    logger = logging.getLogger("VADServer")
    logger.setLevel(logging.INFO)

    # 文件处理器 - 按大小轮转
    file_handler = RotatingFileHandler(
        os.path.join(log_dir, "log.log"),
        maxBytes=5 * 1024 * 1024,  # 10MB
        backupCount=5
    )
    file_handler.setLevel(logging.INFO)

    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # 格式化器
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    # 添加处理器
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    return logger

@dataclass
class AudioSession:
    buffer: deque  # 使用deque替代list来实现固定大小的缓冲区
    last_update: float
    max_buffer_size: int = 10  # 最大保留xx个音频包
    sample_rate: int = 16000  # 默认采样率
    timeout: float = 1.0

    def __init__(self, max_buffer_size: int = 10):
        self.buffer = deque(maxlen=max_buffer_size)
        self.last_update = time.time()
        self.max_buffer_size = max_buffer_size

class VADServer:
    def __init__(self):
        self.logger = logging.getLogger("VADServer")
        self.model = AutoModel(model="fsmn-vad")
        self.sessions: Dict[str, AudioSession] = {}
        self.session_lock = threading.Lock()

        # 启动会话清理线程
        self.cleanup_thread = threading.Thread(target=self._cleanup_sessions, daemon=True)
        self.cleanup_thread.start()

    def process_audio(self, audio_data: bytes, session_id: str) -> Optional[dict]:
        with self.session_lock:
            session = self.sessions.get(session_id)
            if not session:
                self.logger.info(f"{session_id} voice enter")
                session = AudioSession()
                self.sessions[session_id] = session

            # 更新会话 - 自动维护固定大小的缓冲区
            session.buffer.append(audio_data)
            session.last_update = time.time()

            # 将字节数据转换为numpy数组
            try:
                # 只处理缓冲区中的数据
                audio_chunk = np.frombuffer(b''.join(session.buffer), dtype=np.int16)
            except Exception as e:
                self.logger.error(f"Error converting audio data: {e}")
                return None

            # 处理音频
            cache = {}
            result = self.model.generate(
                input=audio_chunk,
                cache=cache,
                is_final=False,
                chunk_size=60  # 60ms chunks
            )

            # 如果检测到语音结束，清除缓存
            if result and result[0].get("value") and self.is_end(result):
                self.logger.info(f"{session_id} voice end {result}")
                session.buffer.clear()  # 清除缓存
                return result[0]

            if result and result[0].get("value"):
                self.logger.info(f"{session_id} voice start {result}")

            return None

    def is_end(self, data):
        # 遍历列表中的每个字典
        for item in data:
            # 检查字典中是否有 'value' 键
            if 'value' in item:
                # 获取 'value' 的值
                value = item['value']
                # 确保 'value' 是一个列表，并且至少有一个子列表
                if isinstance(value, list) and len(value) > 0:
                    # 遍历 'value' 中的每个子列表
                    for sublist in value:
                        # 检查子列表是否是列表，并且长度为2
                        if isinstance(sublist, list) and len(sublist) == 2:
                            # 检查第一个元素是否为 -1，第二个元素是否大于 0
                            if sublist[0] == -1 and sublist[1] > 0:
                                return True
        return False

    def _cleanup_sessions(self):
        while True:
            current_time = time.time()
            with self.session_lock:
                expired_sessions = [
                    session_id for session_id, session in self.sessions.items()
                    if current_time - session.last_update > session.timeout
                ]
                for session_id in expired_sessions:
                    del self.sessions[session_id]
                    self.logger.info(f"delete session {session_id},{len(self.sessions.items())} session remain")
            time.sleep(1)


def main():
    setup_logging()
    context = zmq.Context()
    socket = context.socket(zmq.REP)
    socket.bind("tcp://*:5555")
    # socket.bind("tcp://*:8080")

    vad_server = VADServer()
    vad_server.logger.info("VAD Server started...")

    while True:
        try:
            message_parts = socket.recv_multipart()
            session_id = message_parts[0].decode('utf-8')
            audio_data = message_parts[1]

            # 处理音频数据
            result = vad_server.process_audio(audio_data, session_id)

            # 发送结果
            socket.send_json(result if result else {})

        except Exception as e:
            vad_server.logger.error(f"Error processing request: {e}")
            socket.send_json({"error": str(e)})


if __name__ == "__main__":
    main()